package com.lsy.onehot.rnn.qa;

import org.datavec.api.records.reader.RecordReader;
import org.datavec.api.records.reader.impl.csv.CSVRecordReader;
import org.datavec.api.split.FileSplit;
import org.deeplearning4j.datasets.datavec.RecordReaderDataSetIterator;
import org.deeplearning4j.eval.Evaluation;
import org.deeplearning4j.nn.api.OptimizationAlgorithm;
import org.deeplearning4j.nn.conf.MultiLayerConfiguration;
import org.deeplearning4j.nn.conf.NeuralNetConfiguration;
import org.deeplearning4j.nn.conf.Updater;
import org.deeplearning4j.nn.conf.layers.DenseLayer;
import org.deeplearning4j.nn.conf.layers.OutputLayer;
import org.deeplearning4j.nn.graph.ComputationGraph;
import org.deeplearning4j.nn.multilayer.MultiLayerNetwork;
import org.deeplearning4j.nn.weights.WeightInit;
import org.deeplearning4j.optimize.listeners.ScoreIterationListener;
import org.deeplearning4j.util.ModelSerializer;
import org.nd4j.linalg.api.ndarray.INDArray;
import org.nd4j.linalg.dataset.DataSet;
import org.nd4j.linalg.dataset.api.iterator.DataSetIterator;
import org.nd4j.linalg.lossfunctions.LossFunctions.LossFunction;

import com.gg.lsy.ai.uitls.ini.ZztPropKit;
import com.lsy.lstm.ai.rnn.qa.dataset.io.CorpusHandle;
import com.lsy.onehot.rnn.qa.inner.IOneHotProcess;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.io.LineNumberReader;

/**
 * TODO 替代所有文本处理为数据库处理
 * <p>
 * 
 * @author
 */
public class MultiLayerNetworkOneHotNN {
	private MultiLayerNetwork model;
	private IOneHotProcess ionehotProcess ;
	private File networkFile = new File(CorpusHandle.toTempPath("chat/one-hot-train.zip"));;
	private boolean isLoadNet = false ;
	
	public MultiLayerNetworkOneHotNN()  {
	}
	
	public void setIOneHotProcess(IOneHotProcess ione){
		this.ionehotProcess = ione;
	}

	public MultiLayerNetwork getModel() {
		return model;
	}

	public void initialModel(int numInputs) throws Exception {
		int seed = 123; 
		// 获取结果集数量ICurpusMateSet
		int numOutputs = ionehotProcess.getICurpusMateSet().getDict_A().size() ;
		
		System.out.println("numInputs:"+ numInputs + " numOutputs:"+ numOutputs   );
		boolean isSave = ZztPropKit.getBoolean("network.train.onestep.save", false);
		networkFile = new File(CorpusHandle.toTempPath("chat/one-hot-train.zip"));
		
		System.out.println("network.train.onestep.save:"+isSave + " networkFile:"+ networkFile.getAbsolutePath() );
		//0.01~0.8
		double learningRate = ZztPropKit.getDouble("network.train.learningRate", 0.1 )  ;
//		int offet = ZztPropKit.getInt("network.train.learningRate", 0 )  ;

		int batchSize = 1000;
		int nEpochs = ZztPropKit.getInt("network.train.nEpochs", 500 )  ;
//		int numHiddenNodes = 240;
		int numHiddenNodes = numOutputs + numOutputs/4;
		System.out.println("network.train.learningRate:"+ learningRate ); 
		System.out.println("network.train.nEpochs:"+ nEpochs ); 

		// Load the training data:
		RecordReader rr = new CSVRecordReader();
		// rr.initialize(new FileSplit(new
		// File("src/main/resources/classification/linear_data_train.csv")));  
		//w2v_tr.txt
		rr.initialize(new FileSplit(  ionehotProcess.getOutVec()  ));
		DataSetIterator trainIter = new RecordReaderDataSetIterator(rr, batchSize, 0, numOutputs);

		// Load the test/evaluation data:
		RecordReader rrTest = new CSVRecordReader();
		//w2v_test.txt
		rrTest.initialize(new FileSplit(  ionehotProcess.getOutVec()   ));
		DataSetIterator testIter = new RecordReaderDataSetIterator(rrTest, batchSize, 0, numOutputs);

		// build neural network
		MultiLayerConfiguration conf = new NeuralNetConfiguration.Builder().seed(seed).iterations(1)
				.optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT).regularization(true).l2(0.1)
				.updater(Updater.NESTEROVS).momentum(0.95) // specify the rate
															// of change of the
															// learning rate.
				.weightInit(WeightInit.XAVIER).activation("relu").list()
				.layer(0,
						new DenseLayer.Builder().nIn(numInputs).nOut(numHiddenNodes).learningRate(learningRate).build())
				.layer(1,
						new DenseLayer.Builder().nIn(numHiddenNodes).nOut(numHiddenNodes).learningRate(learningRate)
								.build())
				.layer(2,
						new OutputLayer.Builder(LossFunction.NEGATIVELOGLIKELIHOOD).weightInit(WeightInit.XAVIER)
								.activation("softmax").nIn(numHiddenNodes).nOut(numOutputs).build())
				.pretrain(false).backprop(true).build();

		// UIServer uiServer = UIServer.getInstance();
		// StatsStorage statsStorage = new InMemoryStatsStorage();
		// uiServer.attach(statsStorage);
		
//		if( isLoadNet ){
//			MultiLayerNetwork model = restoreComputationGraph();
//		}
		// 网络训练可视化
		MultiLayerNetwork model;
		if(networkFile.exists() ){
			model = restoreComputationGraph();
			System.out.println("network.train. is load exist:"+ networkFile.getAbsolutePath()  ); 
		}else
		{
			model = new MultiLayerNetwork(conf);
			model.init();
		}
		model.setListeners(new ScoreIterationListener(1));
		// model.setListeners(new StatsListener(statsStorage));

		// Initialize the user interface backend
		for (int n = 0 ; n < nEpochs; n++) {
			model.fit(trainIter);
			//也可以每一步都保存，也可以在后面保存
			if( n % 10 == 0){
				saveModel(networkFile, model );
				System.out.println("nEpochs:"+ nEpochs +" ("+( n*100/nEpochs )+"%)");
			}
		}
		
		saveModel(networkFile, model );
		System.out.println("Evaluate model....");
		Evaluation eval = new Evaluation(numOutputs);
		while (testIter.hasNext()) {
			DataSet t = testIter.next();
			INDArray features = t.getFeatureMatrix();
			INDArray lables = t.getLabels();
			INDArray predicted = model.output(features, false);
			// System.out.println(predicted);
			eval.eval(lables, predicted);
		}
		
		this.model = model;
		saveModel(networkFile, model );
		// Print the evaluation statistics
		// System.out.println(eval.stats());
	}
	
	//restore from serializer model 
	public MultiLayerNetwork restoreComputationGraph(){
		MultiLayerNetwork net = null;
        try {
			net = ModelSerializer.restoreMultiLayerNetwork(networkFile);
	        net.setListeners(new ScoreIterationListener(1));
		} catch (IOException e) {
			// TODO Auto-generated catch block
			e.printStackTrace();
		}
        return net ;
	}
	
	
	private void saveModel(File networkFile,MultiLayerNetwork model) throws IOException {
		System.out.println("Saving the model...");
		File backup = new File(CorpusHandle.toTempPath("chat/one-hot-train-bak.zip"));
		if (networkFile.exists()) {
			if (backup.exists()) {
				backup.delete();
			}
			networkFile.renameTo(backup);
		}
		ModelSerializer.writeModel(model, networkFile, true);
		System.out.println("Done.");
	}
	
}
